/**
 * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved.
 *
 * This file is a part of the CANN Open Software.
 * Licensed under CANN Open Software License Agreement Version 1.0 (the "License").
 * Please refer to the License for details. You may not use this file except in compliance with the License.
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
 * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
 * See LICENSE in the root of the software repository for the full text of the License.
 */

#include <vector>
#include <cstdint>
#include <cmath>
#include <random>
#include <iostream>
#include <algorithm>
#include "acl/acl.h"
#include "atvc.h"

#define CHECK_ACL(x)                                                                        \
    do {                                                                                    \
        aclError __ret = x;                                                                 \
        if (__ret != ACL_ERROR_NONE) {                                                      \
            std::cerr << __FILE__ << ":" << __LINE__ << " aclError:" << __ret << std::endl; \
        }                                                                                   \
    } while (0)

namespace {
    // 描述算子的输入输出以及临时计算资源
    using AxpyOpTraits = ATVC::OpTraits<ATVC::OpInputs<float,float,float>, ATVC::OpOutputs<float>, ATVC::OpTemps<>>;

    // 传入编译态参数ATVC::OpTraits
    template<class Traits>
    // 开发自定义函数名/类名
    struct AxpyComputeFunc {
        // DataType模板参数，根据实际数据类型个数填写
        template<typename T>
        // 重载operator公有接口，提供给`ATVC::Kernel::EleWiseOpTemplate`调用
        __aicore__ inline void operator()(AscendC::LocalTensor<T> a,
                                        AscendC::LocalTensor<T> x,
                                        AscendC::LocalTensor<T> y,
                                        AscendC::LocalTensor<T> z)
        {
            uint32_t tiledCnt = x.GetSize(); 
            // T scalar = a.GetValue(0);              // 从第一个输入 a 中读取标量
            AscendC::Mul(z, a, x, tiledCnt);      // z = a * x
            AscendC::Add(z, z, y, tiledCnt);       // z = z + y = a * x + y
        }
    };
    void InitializeData(int32_t eleNum, std::vector<float> &inputX, std::vector<float> &inputY, float alpha, std::vector<float> &golden)
    {
        std::random_device rd;
        std::mt19937 gen(rd());     
        std::uniform_real_distribution<float> dis(1.0f, 10.0f);
        for (int i = 0; i < eleNum; ++i) {
            inputX[i] = dis(gen);
            inputY[i] = dis(gen);
            // AXPY: y = alpha * x + y
            golden[i] = alpha * inputX[i] + inputY[i];
        }
    }
}

/*
 * 该函数为AxpyCustom算子核函数入口
 * x        Device上的gm地址，指向AxpyCustom算子第一个输入
 * y        Device上的gm地址，指向AxpyCustom算子第一个输出
 * param    Device上的gm地址，指向运行态ATVC::EleWiseParam数据
*/
template<class OpTraits>
__global__ __aicore__ void AxpyCustom(GM_ADDR a, GM_ADDR x, GM_ADDR y, GM_ADDR z, ATVC::EleWiseParam param)
{
    KERNEL_TASK_TYPE_DEFAULT(KERNEL_TYPE_AIV_ONLY); // 控制算子执行时只启动Vector核
    auto op = ATVC::Kernel::EleWiseOpTemplate<AxpyComputeFunc<OpTraits>>();
    op.Run(a, x, y, z, &param);          // 按照输入、输出、param的顺序传入Run函数中；OpTraits内部的ATVC::OpTemps将由EleWiseOpTemplate内部申请资源，开发无需关注
}
bool VerifyResults(const std::vector<float>& golden, const std::vector<float>& actual) {
    if (golden.size() != actual.size()) {
        printf("Size mismatch: golden %zu vs actual %zu\n", golden.size(), actual.size());
        return false;
    }
    const float tolerance = 1e-3f;
    bool hasError = false;
    int errorCount = 0;

    const int maxPrint = 10;
    int printed = 0;
    for (size_t i = 0; i < golden.size(); ++i) {
        float diff = std::abs(golden[i] - actual[i]);
        float relativeError = golden[i] != 0.0f ? diff / std::abs(golden[i]) : diff;

        // 仅打印最多 maxPrint 条比较结果
        if (printed < maxPrint) {
            printf("Index %zu: golden=%f actual=%f diff=%f rel_err=%f\n",
                   i, golden[i], actual[i], diff, relativeError);
            ++printed;
        }

        // 使用相对误差和绝对误差的组合判断错误
        if (diff > tolerance && relativeError > 1e-4f) {
            hasError = true;
            ++errorCount;
        }
    }

    if (printed < static_cast<int>(golden.size())) {
        printf("Displayed %d of %zu elements.\n", printed, golden.size());
    }

    if (hasError) {
        printf("Total verification errors found: %d\n", errorCount);
        return false;
    }

    printf("No verification errors found.\n");
    return true;
}

int main()
{
    // init data
    int32_t eleNum = 8 * 2048;
    float alpha = 7.1f;  // 添加alpha标量值
    size_t scalarByteSize = sizeof(float);
    size_t vectorByteSize = static_cast<size_t>(eleNum) * sizeof(float);
    size_t outputByteSize = static_cast<size_t>(eleNum) * sizeof(float);

    std::vector<float> inputX(eleNum);
    std::vector<float> inputY(eleNum);
    std::vector<float> golden(eleNum);
    InitializeData(eleNum, inputX, inputY, alpha, golden);

    // 新增: 构造与 x/y 等长的 alpha 向量，避免 a 只分配 1 个元素导致越界
    std::vector<float> alphaVec(eleNum, alpha);

    ATVC::EleWiseParam param;

    // 计算输入为8*2048个float元素的Axpy算子的运行态参数param
    if (!ATVC::Host::CalcEleWiseTiling<AxpyOpTraits>(eleNum, param)) {
        printf("Elewise tiling error.");
        return -1;
    }; 
    // 初始化Acl资源与数据 
    aclrtContext context;
    aclrtStream stream = nullptr;
    int32_t deviceId = 0;
    CHECK_ACL(aclInit(nullptr));
    CHECK_ACL(aclrtSetDevice(deviceId));
    CHECK_ACL(aclrtCreateContext(&context, deviceId));
    CHECK_ACL(aclrtCreateStream(&stream));

    // 分配设备内存
    uint8_t *aDevice, *xDevice, *yDevice, *zDevice;
    uint8_t *zHost;
    // CHECK_ACL(aclrtMalloc((void **)&aDevice, scalarByteSize, ACL_MEM_MALLOC_HUGE_FIRST));
    // CHECK_ACL(aclrtMemcpy(aDevice, scalarByteSize, &alpha, scalarByteSize, ACL_MEMCPY_HOST_TO_DEVICE));
    CHECK_ACL(aclrtMalloc((void **)&aDevice, vectorByteSize, ACL_MEM_MALLOC_HUGE_FIRST));
    CHECK_ACL(aclrtMemcpy(aDevice, vectorByteSize, alphaVec.data(), vectorByteSize, ACL_MEMCPY_HOST_TO_DEVICE));
    CHECK_ACL(aclrtMalloc((void **)&xDevice, vectorByteSize, ACL_MEM_MALLOC_HUGE_FIRST));
    CHECK_ACL(aclrtMemcpy(xDevice, vectorByteSize, inputX.data(), vectorByteSize, ACL_MEMCPY_HOST_TO_DEVICE));
    CHECK_ACL(aclrtMalloc((void **)&yDevice, vectorByteSize, ACL_MEM_MALLOC_HUGE_FIRST));
    CHECK_ACL(aclrtMemcpy(yDevice, vectorByteSize, inputY.data(), vectorByteSize, ACL_MEMCPY_HOST_TO_DEVICE));
    CHECK_ACL(aclrtMallocHost((void **)(&zHost), outputByteSize));
    CHECK_ACL(aclrtMalloc((void **)&zDevice, outputByteSize, ACL_MEM_MALLOC_HUGE_FIRST));

    // 调用自定义的Kernel API, <<<>>>的BlockNum参数可通过param的TilingData获取
    AxpyCustom<AxpyOpTraits><<<param.tilingData.blockNum, nullptr, stream>>>(aDevice, xDevice, yDevice, zDevice, param);
    
    CHECK_ACL(aclrtSynchronizeStream(stream));
    CHECK_ACL(aclrtMemcpy(zHost, outputByteSize, zDevice, outputByteSize, ACL_MEMCPY_DEVICE_TO_HOST));
    std::vector<float> outputZ(reinterpret_cast<float*>(zHost), reinterpret_cast<float*>(zHost) + eleNum);

    // 释放ACL资源
    CHECK_ACL(aclrtFree(aDevice));
    CHECK_ACL(aclrtFree(xDevice));
    CHECK_ACL(aclrtFree(yDevice));
    CHECK_ACL(aclrtFree(zDevice));
    CHECK_ACL(aclrtFreeHost(zHost));
    CHECK_ACL(aclrtDestroyStream(stream));
    CHECK_ACL(aclrtDestroyContext(context));
    CHECK_ACL(aclrtResetDevice(deviceId));
    CHECK_ACL(aclFinalize());

    if (!VerifyResults(golden, outputZ)) {
        return -1;
    }
    printf("Accuracy verification passed.\n");
    return 0;
}